"""UNet++数据下载模块"""
import os
import zipfile
import openi
from openi import openi_download_file
from launcher import  get_project_root


download_dir = get_project_root()

# 定义下载函数，包含进度打印
def download_and_unzip_segmentation_datasets():
    openi_download_file("enter/datasets_as_numpy", repo_type="dataset", local_dir=download_dir, max_workers=10)
    print('\n数据集下载已完成')


def download_and_unzip_nested_unet_checkpoints():
    openi_download_file("enter/nested_unet_checkpoints", repo_type="dataset", local_dir=download_dir, max_workers=10)
    print('\n模型下载和解压已完成')


def download_nested_unet_om():
    openi.download_model(repo_id="enter/nodule_segmentation",
                         model_name="nested_unet_om",save_path=download_dir)

    model_path = os.path.join(download_dir, 'nested_unet.om')
    # 检查ZIP文件是否存在
    if os.path.exists(model_path):
        print('模型文件 nested_unet.om 下载已完成')


def download_nested_unet_onnx():
    openi.download_model(repo_id="enter/nodule_segmentation",
                         model_name="nested_unet_onnx",save_path=download_dir)

    model_path = os.path.join(download_dir, 'nested_unet.onnx')
    # 检查ZIP文件是否存在
    if os.path.exists(model_path):
        print('模型文件 nested_unet.onnx 下载已完成')


def download_ultrasound_images():
    openi_download_file("enter/ultrasound_images_to_show", repo_type="dataset", local_dir=os.getcwd(), max_workers=10)
    print('\n超声图片数据下载成功！')


__all__ = ['download_and_unzip_segmentation_datasets',
           'download_and_unzip_nested_unet_checkpoints',
           'download_ultrasound_images',
           'download_nested_unet_om',
           'download_nested_unet_onnx']
